package com.ly.matting.model;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtException;
import ai.onnxruntime.OrtSession;
import com.ly.utils.WHCUtil;
import lombok.SneakyThrows;
import org.opencv.core.*;
import org.opencv.imgcodecs.Imgcodecs;
import org.opencv.imgproc.Imgproc;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.HashMap;
import java.util.Map;

public class EnhancedPortraitMatting {

    private static OrtEnvironment env;
    private static OrtSession session;

    public EnhancedPortraitMatting() {
        initializeOnnxRuntime();
    }

    private void initializeOnnxRuntime() {
        env = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions options = new OrtSession.SessionOptions();
        try (InputStream is = getClass().getResourceAsStream("/onnx/rvm_mobilenetv3_fp32.onnx")) {
            if (is == null) {
                throw new RuntimeException("rvm_mobilenetv3_fp32.onnx not found");
            }
            byte[] modelBytes = readAllBytes(is);
            session = env.createSession(modelBytes, options);
        } catch (IOException | OrtException e) {
            throw new RuntimeException("Could not create session", e);
        }
        nu.pattern.OpenCV.loadLocally();
    }

    private byte[] readAllBytes(InputStream is) throws IOException {
        ByteArrayOutputStream buffer = new ByteArrayOutputStream();
        int nRead;
        byte[] data = new byte[16384];
        while ((nRead = is.read(data, 0, data.length)) != -1) {
            buffer.write(data, 0, nRead);
        }
        return buffer.toByteArray();
    }

    @SneakyThrows
    public float[] portraitMatting(float[] imgData, float[] bgData, int width, int height) {
        OnnxTensor inputTensor = createInputTensor(imgData, width, height);
        Map<String, OnnxTensor> inputs = prepareInputs(inputTensor);

        try (OrtSession.Result results = session.run(inputs)) {
            float[] phaData = ((OnnxTensor) results.get(1)).getFloatBuffer().array();
            return applyMatting(imgData, bgData, phaData, width, height);
        } finally {
            releaseInputs(inputs);
        }
    }

    private float[] applyMatting(float[] imgData, float[] bgData, float[] phaData, int width, int height) {
        float[] result = new float[width * height * 4];
        for (int i = 0; i < width * height; i++) {
            float alpha = phaData[i];
            for (int c = 0; c < 3; c++) {
                result[i * 4 + c] = imgData[i * 3 + c] * alpha + bgData[i * 3 + c] * (1 - alpha);
            }
            result[i * 4 + 3] = 1.0f; // 设置完全不透明
        }
        return result;
    }

    private OnnxTensor createInputTensor(float[] imgData, int width, int height) throws OrtException {
        float[] chw = WHCUtil.whc2cwh(imgData);
        long[] srcShape = new long[]{1, 3, height, width};
        return OnnxTensor.createTensor(env, FloatBuffer.wrap(chw), srcShape);
    }

    private Map<String, OnnxTensor> prepareInputs(OnnxTensor inputTensor) throws OrtException {
        Map<String, OnnxTensor> inputs = new HashMap<>();
        inputs.put("src", inputTensor);
        inputs.put("r1i", OnnxTensor.createTensor(env, new float[1][1][1][1]));
        inputs.put("r2i", OnnxTensor.createTensor(env, new float[1][1][1][1]));
        inputs.put("r3i", OnnxTensor.createTensor(env, new float[1][1][1][1]));
        inputs.put("r4i", OnnxTensor.createTensor(env, new float[1][1][1][1]));
        inputs.put("downsample_ratio", OnnxTensor.createTensor(env, new float[]{0.25f}));
        return inputs;
    }

    private void releaseInputs(Map<String, OnnxTensor> inputs) {
        for (OnnxTensor tensor : inputs.values()) {
            tensor.close();
        }
    }


    public static void main(String[] args) {
        long l = System.currentTimeMillis();
        EnhancedPortraitMatting mattingModel = new EnhancedPortraitMatting();

        // 读取前景图像
        Mat img = Imgcodecs.imread("1.png");

        // 读取背景图像
//        Mat bg = Imgcodecs.imread("background.png");
        Mat bg = new Mat(1080, 1920, CvType.CV_8UC3, new Scalar(0, 255, 0));

        // 确保图像是 3 通道的
        if (img.channels() != 3 || bg.channels() != 3) {
            throw new RuntimeException("Input images must have 3 channels");
        }

        // 确保背景图像与前景图像大小相同
        if (!img.size().equals(bg.size())) {
            Imgproc.resize(bg, bg, img.size());
        }

        // 转换图像为浮点型
        Mat imgFloat = new Mat();
        Mat bgFloat = new Mat();
        img.convertTo(imgFloat, CvType.CV_32FC3, 1.0/255.0);
        bg.convertTo(bgFloat, CvType.CV_32FC3, 1.0/255.0);

        // 获取图像数据
        float[] imgData = new float[img.rows() * img.cols() * 3];
        float[] bgData = new float[bg.rows() * bg.cols() * 3];
        imgFloat.get(0, 0, imgData);
        bgFloat.get(0, 0, bgData);

        // 执行抠图操作
        float[] result = mattingModel.portraitMatting(imgData, bgData, img.cols(), img.rows());

        // 创建结果 Mat
        Mat resultMat = new Mat(img.rows(), img.cols(), CvType.CV_32FC4);
        resultMat.put(0, 0, result);

        // 转换结果为 8 位无符号整型
        Mat resultMat8U = new Mat();
        resultMat.convertTo(resultMat8U, CvType.CV_8UC4, 255.0);
        System.out.println(System.currentTimeMillis() - l);
        // 保存结果
        Imgcodecs.imwrite("output_with_background.png", resultMat8U);
    }
}